package edu.corn.statemachine.lock;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 采用乐观锁的形式向状态机提供锁机制。
 * 可能会发生死锁。
 */
public class HashMapLockPool implements LockPool {

    private final ConcurrentHashMap<Object, AtomicInteger> lockPool = new ConcurrentHashMap<>();
    private static final int MAX_RETRIES = 10;

    @Override
    public void tryLock(Object key) {
        int currentThreadId = (int) Thread.currentThread().getId();
        if (key == null) key = currentThreadId;
        int retries = 0;

        while (retries < MAX_RETRIES) {
            lockPool.compute(key, (k, existingLock) -> {
                if (existingLock == null) {
                    // 如果不存在锁，则初始化为当前线程 ID
                    return new AtomicInteger(currentThreadId);
                } else {
                    int existingThreadId = existingLock.get();
                    // 尝试更新为当前线程 ID，使用乐观锁
                    if (existingThreadId == 0 && existingLock.compareAndSet(0, currentThreadId)) {
                        return existingLock;
                    } else if (existingThreadId == currentThreadId) {
                        // 已经是当前线程持有的锁
                        return existingLock;
                    }
                }
                return existingLock; // 保持原来的锁状态
            });

            if (isLockedByCurrentThread(key)) {
                return; // 成功获取锁
            }

            retries++;
            try {
                Thread.sleep(10); // 小小的等待，避免过于频繁的重试
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }

    @Override
    public void unlock(Object key) {
        int currentThreadId = (int) Thread.currentThread().getId();
        if (key == null) key = currentThreadId;

        lockPool.computeIfPresent(key, (k, existingLock) -> {
            if (existingLock.get() == currentThreadId) {
                // 只有持有锁的线程可以解锁
                existingLock.set(0);
            }
            return existingLock;
        });
    }

    // 检查锁是否被当前线程持有
    public boolean isLockedByCurrentThread(Object key) {
        int currentThreadId = (int) Thread.currentThread().getId();
        if (key == null) key = currentThreadId;
        return lockPool.getOrDefault(key, new AtomicInteger(0)).get() == currentThreadId;
    }
}
